import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import load_wine
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.tree import DecisionTreeClassifier, plot_tree
from sklearn.metrics import classification_report, confusion_matrix, ConfusionMatrixDisplay

# 设置中文字体（解决中文显示问题）
plt.rcParams['font.sans-serif'] = ['SimHei']  # 微软雅黑
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# ======================
# 1. 数据准备
# ======================
# 加载葡萄酒数据集（sklearn内置数据集）
wine = load_wine()
X = wine.data  # 获取特征数据 (178个样本，13个特征)
y = wine.target  # 获取类别标签 (3个类别)
feature_names = [
    '酒精浓度', '苹果酸', '灰分', '灰分碱度', '镁含量',
    '总酚', '类黄酮', '非黄酮酚', '原花青素',
    '颜色强度', '色调', 'OD280/OD315', '脯氨酸'
]  # 自定义中文特征名称

# 划分训练集和测试集 (8:2比例)
X_train, X_test, y_train, y_test = train_test_split(
    X, y,
    test_size=0.2,
    random_state=42  # 固定随机种子确保可复现性
)

# ======================
# 2. 模型参数优化
# ======================
# 定义参数搜索空间
params = {
    'max_depth': [3, 5, 7, None],  # 树的最大深度
    'min_samples_split': [2, 5, 10],  # 节点分裂最小样本数
    'min_samples_leaf': [1, 2, 4],  # 叶节点最小样本数
    'criterion': ['gini', 'entropy']  # 分裂标准：基尼系数/信息熵
}

# 创建网格搜索对象
grid = GridSearchCV(
    estimator=DecisionTreeClassifier(),  # 基础模型
    param_grid=params,  # 参数空间
    cv=5,  # 5折交叉验证
    n_jobs=-1  # 使用全部CPU核心
)

# 执行网格搜索
grid.fit(X_train, y_train)

# 输出最佳参数
best_tree = grid.best_estimator_
print("最佳参数组合:", grid.best_params_)
print("训练集准确率: {:.2f}".format(best_tree.score(X_train, y_train)))
print("测试集准确率: {:.2f}".format(best_tree.score(X_test, y_test)))

# ======================
# 3. 模型评估与可视化
# 生成预测结果
y_pred = best_tree.predict(X_test)

# 打印分类报告
print("\n分类评估报告:")
print(classification_report(y_test, y_pred, target_names=wine.target_names))

# ----------------------
# 可视化1：决策树结构
# ----------------------
plt.figure(figsize=(14, 7))
plot_tree(
    decision_tree=best_tree,
    feature_names=feature_names,  # 使用中文特征名
    class_names=wine.target_names,  # 类别名称
    filled=True,  # 颜色填充
    rounded=True,  # 圆角节点
    fontsize=10,  # 字体大小
    proportion=True  # 显示样本比例
)
plt.title("优化后的决策树结构", fontsize=18)
plt.show()

# ----------------------
# 可视化2：特征重要性
# ----------------------
plt.figure(figsize=(12, 6))
importances = best_tree.feature_importances_
indices = np.argsort(importances)[::-1]  # 按重要性降序排列

plt.bar(
    range(len(importances)),
    importances[indices],
    align='center',
    color='#1f77b4'  # 统一颜色
)

# 设置坐标轴标签
plt.xticks(range(len(importances)), np.array(feature_names)[indices], rotation=45, ha='right')
plt.xlabel("特征名称", fontsize=12)
plt.ylabel("重要性得分", fontsize=12)
plt.title("特征重要性排序", fontsize=16)
plt.tight_layout()  # 自动调整布局
plt.show()

# ----------------------
# 可视化3：混淆矩阵
# ----------------------
cm = confusion_matrix(y_test, y_pred)
disp = ConfusionMatrixDisplay(
    confusion_matrix=cm,
    display_labels=wine.target_names
)
disp.plot(
    cmap=plt.cm.Blues,  # 颜色映射
    values_format='d',  # 显示整数
    colorbar=False  # 隐藏颜色条
)
plt.title("混淆矩阵", fontsize=16)
plt.grid(False)  # 去除网格线
plt.show()


# ----------------------
# 可视化4：决策边界（前两个特征）
# ----------------------
def plot_decision_surface():
    # 取前两个特征简化可视化
    X_2d = X[:, :2]

    # 训练简化模型
    tree_2d = DecisionTreeClassifier(**grid.best_params_)
    tree_2d.fit(X_2d, y)

    # 生成网格数据
    x_min, x_max = X_2d[:, 0].min() - 1, X_2d[:, 0].max() + 1
    y_min, y_max = X_2d[:, 1].min() - 1, X_2d[:, 1].max() + 1
    xx, yy = np.meshgrid(
        np.arange(x_min, x_max, 0.02),
        np.arange(y_min, y_max, 0.02)
    )

    # 预测网格点类别
    Z = tree_2d.predict(np.c_[xx.ravel(), yy.ravel()])
    Z = Z.reshape(xx.shape)

    # 绘制决策边界
    plt.figure(figsize=(12, 8))
    plt.contourf(xx, yy, Z, alpha=0.4, cmap=plt.cm.Paired)

    # 绘制原始数据点
    scatter = plt.scatter(
        X_2d[:, 0], X_2d[:, 1],
        c=y,
        edgecolors='k',
        cmap=plt.cm.Paired
    )

    # 添加图例和标签
    legend = plt.legend(
        *scatter.legend_elements(),
        title="类别",
        loc="upper right"
    )
    plt.xlabel(feature_names[0], fontsize=12)
    plt.ylabel(feature_names[1], fontsize=12)
    plt.title("基于前两个特征的决策边界可视化", fontsize=16)
    plt.show()


plot_decision_surface()